Skip to content

Support DoRA (use_dora=True) in the safetensors LoRA merge#828

Open
danielhanchen wants to merge 1 commit into
unslothai:mainfrom
danielhanchen:dora-merge-support
Open

Support DoRA (use_dora=True) in the safetensors LoRA merge#828
danielhanchen wants to merge 1 commit into
unslothai:mainfrom
danielhanchen:dora-merge-support

Conversation

@danielhanchen

Copy link
Copy Markdown
Member

Summary

DoRA (use_dora=True) adapters could not be merged through the safetensors path. create_lora_statistics did not capture the lora_magnitude_vector, so the key-consistency check raised Unsloth: Extracted keys = {...lora_magnitude_vector.default.weight} do not match!, and even past that the merge math applied only W0 + alpha*BA, dropping the magnitude.

This adds DoRA support for the dense merge path (attention / MLP / any dense linear, including the language and vision/audio towers of multimodal models):

  • LoraStats gains a magnitude field; create_lora_statistics captures lora_magnitude_vector.default.
  • assert_same_keys ignores the magnitude tensor (folded into the weight, omitted from the merged checkpoint), like lora_A / lora_B.
  • _merge_lora folds it in as (m / ||W0 + alpha*BA||_row) * (W0 + alpha*BA), i.e. one L2 norm per output row over the input dim, matching PEFT's DoRA merge.

MoE-expert DoRA is explicitly refused with a clear error (the per-expert / fused expert helpers fold only the LoRA delta, not the magnitude), so a DoRA MoE merge fails loud instead of silently dropping the magnitude.

Test

tests/test_dora_merge.py:

  • dense DoRA merge equals PEFT's own merge_and_unload (atol 1e-4) and actually changes the weight;
  • plain LoRA (use_dora=False) is unchanged (magnitude is None);
  • DoRA on a MoE expert raises.

Full merge e2e suite (dense + MoE fused/per-expert + resized) stays green (65 passed, 6 skipped).

The dense merge ignored the DoRA magnitude vector: create_lora_statistics did not capture
it, so the key-consistency check raised 'Extracted keys do not match', and the merge math
applied only W0 + alpha*BA. Now capture lora_magnitude_vector, ignore it in the key check,
and fold it in _merge_lora as (m / ||W0 + alpha*BA||_row) * (W0 + alpha*BA), matching PEFT's
DoRA merge. MoE-expert DoRA is refused (fail loud) since the expert helpers do not apply the
magnitude. Test compares the dense merge against PEFT merge_and_unload.
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request adds support for merging DoRA (use_dora=True) adapters in the safetensors merge path by rescaling the merged weight direction to the learned magnitude vector. It also ensures that DoRA magnitude weights are correctly ignored during key-consistency checks and explicitly refuses DoRA merges on MoE expert layers to prevent silent failures. A comprehensive test suite is introduced to validate these changes against PEFT's own DoRA merge. No review comments were provided, so there is no additional feedback.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant